import os
import random
import paddle
import paddle.fluid as fluid
from paddle.fluid.dygraph.nn import Conv2D, Pool2D, Linear
import numpy as np
from PIL import Image
import gzip
import json


# 定义数据集读取器
def load_data():

    # 数据文件
    datafile = './data/data49109/mnist.json.gz'
    print('mnist数据集从 {} 加载中  ......'.format(datafile))
    # 打开数据中的数据并转为JSON格式数据
    data = json.load(gzip.open(datafile))
    # 将JSON格式数据 分为 训练数据集 ,验证数据集,观察集
    train_set, val_set, eval_set = data

    # 数据集相关参数，图片高度IMG_ROWS, 图片宽度IMG_COLS
    IMG_ROWS = 28
    IMG_COLS = 28
    # 直接使用观察集
    imgs = eval_set[0]
    labels = eval_set[1]

    # 获得所有图像的数量
    imgs_length = len(imgs)

    # 验证图像数量和标签数量是否一致
    assert len(imgs) == len(labels), \
          "length of train_imgs({}) should be the same as train_labels({})".format(
                  len(imgs), len(labels))
    # 获得索引列表
    index_list = list(range(imgs_length))

    # 读入数据时用到的batchsize
    BATCHSIZE = 100
    random.shuffle(index_list)

    # 定义数据生成器
    def data_generator():
        imgs_list = []
        labels_list = []
        for i in index_list:
            img = np.reshape(imgs[i],
                             [1, IMG_ROWS, IMG_COLS]).astype('float32')
            label = np.reshape(labels[i], [1]).astype('int64')
            imgs_list.append(img)
            labels_list.append(label)
            if len(imgs_list) == BATCHSIZE:
                yield np.array(imgs_list), np.array(labels_list)
                imgs_list = []
                labels_list = []
                break

        # 如果剩余数据的数目小于BATCHSIZE，
        # 则剩余数据一起构成一个大小为len(imgs_list)的mini-batch
        if len(imgs_list) > 0:
            yield np.array(imgs_list), np.array(labels_list)

    return data_generator


# 定义模型结构
class MNIST(fluid.dygraph.Layer):
    def __init__(self):
        super(MNIST, self).__init__()

        # 定义一个卷积层，使用relu激活函数
        self.conv1 = Conv2D(num_channels=1,
                            num_filters=20,
                            filter_size=5,
                            stride=1,
                            padding=2,
                            act='relu')
        # 定义一个池化层，池化核为2，步长为2，使用最大池化方式
        self.pool1 = Pool2D(pool_size=2, pool_stride=2, pool_type='max')
        # 定义一个卷积层，使用relu激活函数
        self.conv2 = Conv2D(num_channels=20,
                            num_filters=20,
                            filter_size=5,
                            stride=1,
                            padding=2,
                            act='relu')
        # 定义一个池化层，池化核为2，步长为2，使用最大池化方式
        self.pool2 = Pool2D(pool_size=2, pool_stride=2, pool_type='max')
        # 定义一个全连接层，输出节点数为10
        self.fc = Linear(input_dim=980, output_dim=10, act='softmax')
# 定义网络的前向计算过程

    def forward(self, inputs):
        x = self.conv1(inputs)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.pool2(x)
        x = fluid.layers.reshape(x, [x.shape[0], 980])
        x = self.fc(x)
        return x


# 读取数据
eval_data = load_data()

#在使用GPU机器时，可以将use_gpu变量设置成True
use_gpu = True
place = fluid.CUDAPlace(0) if use_gpu else fluid.CPUPlace()

# 正确的比值
correctRate = 0
err = []
# 定义预测过程
with fluid.dygraph.guard(place):
    model = MNIST()
    params_file_path = 'mnist9'

    # 加载模型参数
    model_dict, _ = fluid.load_dygraph("mnist9")
    model.load_dict(model_dict)
    model.eval()
    for datas in eval_data():
        imgs, labels = datas
        imgs = imgs.reshape(100, 1, 1, 28, 28)
        for index in range(len(imgs)):
            # print(index)
            results = model(fluid.dygraph.to_variable(imgs[index]))
            lab = np.argsort(results.numpy())
            print("第{}次预测的数字是{},正确答案是{}: ".format(index, lab[0][-1],
                                                  labels[index][0]))
            if lab[0][-1] == labels[index][0]:
                correctRate = correctRate + 1
            else:
                err.append(index)
print('本模型的识别率是：{}%'.format(correctRate / len(labels) * 100))
print('错误的行数是：{}'.format(err))